library(maptools)
library(tidyverse)
library(igraph)

## ----------------
## 250-precinct map
## ----------------
fl_map <- readShapePoly("map_data/fl/FL.shp")

## Load nodes
nodes <- strsplit(readLines("map_data/map_sub_nodes_250.dat")," ")[[1]]

## First, subset down to right map
fl_sub <- fl_map[which(rownames(fl_map@data) %in% nodes),]

## Then, reorder nodes to be correct order to match - otherwise, won't
## line up properly with original map
fl_sub <- fl_sub[match(nodes, rownames(fl_sub@data)),]

## Convert factor columns to characters, because this particular data is messy
indx <- sapply(fl_sub@data, is.factor)
fl_sub@data[indx] <- lapply(
    fl_sub@data[indx], function(x) as.numeric(as.character(x))
)

## Load edgelist representation of map
el <- strsplit(readLines("map_data/map_sub_250.dat"), split = " ")
el <- apply(do.call(rbind, el), 2, as.numeric)

## Load 1000-solution subset of solutions
## There are 100-million solutions total across all files
sols_all <- readLines("/n/imai_lab/bfifield/enumeration/largemap_enum_sample/solutions_1.dat", n = 1000)

## Convert solutions to congressional district membership matrix
ndists <- 2
sols_out_all <- lapply(1:length(sols_all), function(x){
    sol_split <- as.numeric(strsplit(sols_all[x], split = " ")[[1]])
    el_sub <- el[sol_split,]
    comps_out <- components(graph_from_edgelist(el_sub, directed = FALSE))
    if(comps_out$no < ndists){
        max_num <- max(comps_out$membership)
        comps_out <- c(comps_out$membership, (max_num + 1):ndists)
    }else{
        comps_out <- comps_out$membership
    }
    return(comps_out)
})
sols_out_all <- do.call(cbind, sols_out_all)

## --------
## Iowa map
## --------
## Load map and join in auxiliary data
ia <- readShapePoly("map_data/ia/county.shp")
votes <- read_csv("map_data/ia/ia_vote.csv")
pops <- read_csv("map_data/ia/ia_pop.csv") %>%
    mutate(County = gsub(" County", "", County))
info <- inner_join(votes, pops) %>%
    rename(county = County,
           trump = `Trump #`,
           clinton = `Clinton #`,
           other = `Others #`,
           total = `Total`,
           fips = `FIPS code[10]`,
           pop = Population) %>%
    select(county, trump, clinton, fips, pop)

## Join auxiliary data with shapefile
ia@data$COUNTY <- as.character(ia@data$COUNTY)
ia@data$COUNTY[ia@data$COUNTY == "Obrien"] <- "O'Brien"
ia@data$id <- 1:nrow(ia@data)
ia@data <- inner_join(ia@data, info, by = c("COUNTY" = "county"))

## Load and convert to matrix (this one is simpler, since not in
## edgelist form). Note this is 1000 solutions, there are
## 500 million total
sols_all <- readLines("/n/imai_lab/bfifield/enumeration/ia_enumerate/ia_sols_1.dat", n = 1000)
sols_all <- apply(do.call("cbind", strsplit(sols_all, " ")), 2, as.numeric)
